import torch
from torch import nn
from torch.nn import functional as F
from .gw import GWLoss
import numpy as np
import cv2
from rainbowneko.train.loss import CharbonnierLoss, MSSSIMLoss

class CDCLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.gw_loss = GWLoss()
        self.cb_loss = CharbonnierLoss(size_average=False)
        self.ssim_loss = MSSSIMLoss()

    def get_coe(self, hr):
        hr_Y = ((0.257 * hr[:, :1, :, :] + 0.564 * hr[:, 1:2, :, :] + 0.098 * hr[:, 2:, :, :] + 16 / 255.0) * 255.0)
        map_corner = torch.zeros_like(hr_Y)
        map_edge = torch.zeros_like(hr_Y)
        map_flat = torch.zeros_like(hr_Y)

        hr_Y_numpy = np.transpose(hr_Y.cpu().float().numpy(), (0, 2, 3, 1))
        for i in range(hr_Y_numpy.shape[0]):
            dst = cv2.cornerHarris(hr_Y_numpy[i, :, :, 0], 3, 3, 0.04)
            thres1 = 0.01 * dst.max()
            thres2 = -0.001 * dst.max()
            map_corner[i, :, :, :] = torch.from_numpy(np.float32(dst > thres1)).to(hr.device, dtype=hr.dtype)
            map_edge[i, :, :, :] = torch.from_numpy(np.float32(dst < thres2)).to(hr.device, dtype=hr.dtype)
            map_flat[i, :, :, :] = torch.from_numpy(np.float32((dst > thres2) & (dst < thres1))).to(hr.device, dtype=hr.dtype)

        return map_flat, map_edge, map_corner

    def forward(self, sr_list, hr):
        hr = hr.float()
        coe_list = self.get_coe(hr)
        loss = 0
        for i in range(len(sr_list)):
            if i == len(sr_list) - 1:
                loss += self.gw_loss(sr_list[i], hr) + self.ssim_loss(sr_list[i], hr)
            else:
                coe = coe_list[i]
                loss += (self.cb_loss(coe * sr_list[i], coe * hr).sum()/coe.sum())

        return loss
    
    def to(self, *args, **kwargs):
        print(args, kwargs)
        return super().to(*args, **kwargs)
